package com.summerbird.mallchat.common.chat.service.strategy.frequency;

import com.summerbird.mallchat.common.common.service.LockService;
import com.summerbird.mallchat.common.common.utils.RedisUtils;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @description: 令牌桶限流算法
 * @author SummerBird
 * @date 2024/1/30 18:52
 * @version 1.0
 */
@Component
@Slf4j
public class TokenBucketFrequencyStrategy extends AbstractFrequencyControlStrategy<FrequencyControlDTO> {

    @Autowired
    private LockService lockService;

    @Autowired
    private RedissonClient redissonClient;

    private static final int TOKEN_TAKE = 1;

    private static final Map<TimeUnit, RateIntervalUnit> TIME_UNIT_RATE_INTERVAL_UNIT_MAP = new ConcurrentHashMap<>();

    static {
        TIME_UNIT_RATE_INTERVAL_UNIT_MAP.put(TimeUnit.MILLISECONDS, RateIntervalUnit.MILLISECONDS);
        TIME_UNIT_RATE_INTERVAL_UNIT_MAP.put(TimeUnit.SECONDS, RateIntervalUnit.SECONDS);
        TIME_UNIT_RATE_INTERVAL_UNIT_MAP.put(TimeUnit.MINUTES, RateIntervalUnit.MINUTES);
        TIME_UNIT_RATE_INTERVAL_UNIT_MAP.put(TimeUnit.HOURS, RateIntervalUnit.HOURS);
        TIME_UNIT_RATE_INTERVAL_UNIT_MAP.put(TimeUnit.DAYS, RateIntervalUnit.DAYS);
    }

    @Override
    protected boolean reachRateLimit(Map<String, FrequencyControlDTO> frequencyControlMap) {
        ArrayList<String> keyList = new ArrayList<>(frequencyControlMap.keySet());
        for (int i = 0; i < keyList.size(); i++) {
            String key = keyList.get(i);
            RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
            FrequencyControlDTO frequencyControlDTO = frequencyControlMap.get(key);
            rateLimiter.trySetRate(RateType.OVERALL, frequencyControlDTO.count, frequencyControlDTO.time, TIME_UNIT_RATE_INTERVAL_UNIT_MAP.get(frequencyControlDTO.unit));
            // 尝试从令牌桶中获取令牌
            // 如果获取失败，则表示限流了
            if(!rateLimiter.tryAcquire(TOKEN_TAKE)){
                //频率超过了
                log.warn("frequencyControl limit key:{},count:{}", key, frequencyControlDTO.count);
                return true;
            }
        }
        return false;
    }

    @Override
    protected void addFrequencyControlStatisticsCount(Map<String, FrequencyControlDTO> frequencyControlMap) {
        // 令牌消耗在reachRateLimit就已经完成了
        // 这里设置一下限流器的时间即可，超过一天则将限流器删除
        frequencyControlMap.forEach((k,v) -> {
            RedisUtils.expire(k, 1, TimeUnit.DAYS);
        });
    }

    @Override
    protected String getStrategyName() {
        return FrequencyStrategyType.TOKENBUCKET.getName();
    }
}
